"""Process the DOSCAR file made by Vasp
"""
import os
import numpy as np
from .vaspprocess import VaspProcess, log
from .monkeypatch import monkeypatch_class

DOS_DICT = {'nions': 0, 'emax': 1, 'emin': 2, 'nedos': 3, 'efermi': 4}

@monkeypatch_class(VaspProcess)
def initDOS(self):
    log.debug('initDOS')

    if(self.dos_properties is None):
        log.debug('Initializing...')
        try:
            with open(os.path.join(self.directory, 'DOSCAR'),'r') as f:
                line=f.readline()
                dos_properties = [int(line.split()[1])]
                for __ in range(0, 4):
                    f.readline()
                line=f.readline()
                dos_properties += (float(line.split()[0]), float(line.split()[1]), int(line.split()[2]), float(line.split()[3]))
                self.dos_properties=tuple(dos_properties)
                log.debug('{\n'
                          '    NIons    = 'f'{self.dos_properties[DOS_DICT["nions"]]}''\n'
                          '    Emax     = 'f'{self.dos_properties[DOS_DICT["emax"]]}'' eV\n'
                          '    Emin     = 'f'{self.dos_properties[DOS_DICT["emin"]]}'' eV\n'
                          '    NEDOS    = 'f'{self.dos_properties[DOS_DICT["nedos"]]}''\n'
                          '    Efermi   = 'f'{self.dos_properties[DOS_DICT["efermi"]]}'' eV\n'
                          '}')

        except IOError as e:
            log.error('File DOSCAR does not exist or permission denied!')
            raise e


@monkeypatch_class(VaspProcess)
def getDOS(self):
    log.debug('getDOS')

    if(self.dos_properties is None):
        self.initDOS()
    
    energies, DOS = [], []

    try:
        with open(os.path.join(self.directory, 'DOSCAR'),'r') as f:
            for __ in range(0, 6):
                f.readline()
            
            for __ in range(0, self.dos_properties[DOS_DICT['nedos']]):
                line = f.readline()
                e, d, intd = line.split()
                energies.append(float(e))
                DOS.append(float(d))

        return energies, DOS

    except IOError as e:
        log.error('File DOSCAR does not exist or permission denied!')
        raise e

@monkeypatch_class(VaspProcess)
def getPDOS(self, atom: int, orbit: str):
    log.debug('getDOS')

    if(self.dos_properties is None):
        self.initDOS()

    energies, PDOS = [], []

    if(atom >= self.dos_properties[DOS_DICT['nions']]):
        log.error('Atom index out of range!')
        raise IndexError

    ORBIT_DICT = {'s': 1, 
                  'p': 2, 
                  'd': 3, 
                  'p_y': 2, 
                  'p_z': 3,
                  'p_x': 4,
                  'd_{xy}': 5, 
                  'd_{yz}': 6, 
                  'd_{z2-r2}': 7, 
                  'd_{xz}': 8, 
                  'd_{x2-y2}': 9,
                  }

    flag = False

    if(orbit not in ['p', 'd']):
        flag=True

    if(orbit not in ORBIT_DICT):
        log.error('Orbit must be one of the '+str(list(ORBIT_DICT.keys())))
        raise KeyError

    orbit = ORBIT_DICT[orbit]

    try:
        with open(os.path.join(self.directory, 'DOSCAR'),'r') as f:
            for __ in range(0, 6):
                f.readline()

            for __ in range(0, self.dos_properties[DOS_DICT['nedos']]):
                f.readline()

            for __ in range(0, atom):
                if(f.readline() == ''):
                    log.error('There are no PDOS data in DOSCAR file, make sure you have generate them by LORBIT!')
                    raise RuntimeError
                for __ in range(0, self.dos_properties[DOS_DICT['nedos']]):
                    f.readline()

            f.readline()
            for __ in range(0, self.dos_properties[DOS_DICT['nedos']]):
                line = f.readline()
                lines = line.split()

                if(orbit != 1):
                    if((flag == False and len(lines) != 4) or (flag == True and len(lines) != 10)):
                        log.error('It seems something went wrong, check out the DOSCAR file and make sure you properly set LORBIT!')
                        raise RuntimeError

                e = lines[0]
                d = lines[orbit]

                energies.append(float(e))
                PDOS.append(float(d))

        return energies, PDOS

    except IOError as e:
        log.error('File DOSCAR does not exist or permission denied!')
        raise e
